深入理解TVM:内存分配器
一、从make_object开始
TVM的内存分配器用来给前文Object家族中介绍的Object系列类分配内存和构造对象,同时提供了make_object这个helper function来完成这个工作,make_object在TVM中用的非常多,随便举个例子来说,比如src/te/tensor.cc中Tensor类的一个构造函数中就通过make_object来构造了一个TensorNode的对象:
Tensor::Tensor(Array<PrimExpr> shape, DataType dtype, Operation op,
int value_index) {
auto n = make_object<TensorNode>();
n->shape = std::move(shape);
n->dtype = dtype;
n->op = op;
n->value_index = value_index;
data_ = std::move(n);
}
Tensor类继承自DataProducer,DataProducer继承自ObjectRef,上面Tensor构造函数中make_object构造的对象在完成相关初始化之后,最终交由data_管理,如下面代码所示,data_定义在Tensor的顶层基类ObjectRef中,可以看作是一个指向Object对象的指针,Object、ObjectPtr、ObjectRef三者的关系已经在前文Object家族中讲过了:
class ObjectRef {
protected:
ObjectPtr<Object> data_;
};
前面是一个make_object这个helper function使用的具体例子,下面来看make_object具体的定义,在include/tvm/runtime/memory.h中:
template <typename T, typename... Args>
inline ObjectPtr<T> make_object(Args&&... args) {
return SimpleObjAllocator().make_object<T>(std::forward<Args>(args)...);
}
可以看到,make_object内部调用了SimpleObjAllocator的相关函数,SimpleObjAllocator就是本文要说的内存分配器。
make_object只是构造单个的Object,其实除了make_object之外,还有另外一个helper function make_inplace_array_object用于构造Object数组,同样是使用了SimpleObjAllocator这个内存分配器,下面是它的定义:
template <typename ArrayType, typename ElemType, typename... Args>
inline ObjectPtr<ArrayType> make_inplace_array_object(size_t num_elems, Args&&... args) {
return SimpleObjAllocator().make_inplace_array<ArrayType, ElemType>(
num_elems, std::forward<Args>(args)...);
}
下面详细分析下SimpleObjAllocator这个内存分配器。
二、ObjAllocatorBase
SimpleObjAllocator继承自ObjAllocatorBase,两者的类关系定义如下:
template <typename Derived> class ObjAllocatorBase {};
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {};
这里用到了C++的一种编程技巧CRTP,它既可以实现静态多态,又可以复用代码,CRTP在TVM有多处应用,其它的如AttrRegistry也用到了CRTP,这里不详细说了,以后抽空单独写篇文章来详细分析CRTP。
先看基类ObjAllocatorBase,它有两个成员函数,一个是构造单个Object的make_object,一个是构造Object数组的make_inplace_array:
template <typename Derived>
class ObjAllocatorBase {
template <typename T, typename... Args>
ObjectPtr<T> make_object(Args&&... args) {
using Handler = typename Derived::template Handler<T>;
T* ptr = Handler::New(static_cast<Derived*>(this), std::forward<Args>(args)...);
ptr->type_index_ = T::RuntimeTypeIndex();
ptr->deleter_ = Handler::Deleter();
return ObjectPtr<T>(ptr);
}
template <typename ArrayType, typename ElemType, typename... Args>
ObjectPtr<ArrayType> make_inplace_array(size_t num_elems, Args&&... args) {
using Handler = typename Derived::template ArrayHandler<ArrayType, ElemType>;
ArrayType* ptr = Handler::New(static_cast<Derived*>(this), num_elems, std::forward<Args>(args)...);
ptr->type_index_ = ArrayType::RuntimeTypeIndex();
ptr->deleter_ = Handler::Deleter();
return ObjectPtr<ArrayType>(ptr);
}
};
从上面的定义中可以看出,make_object和make_inplace_array的处理流程相同,都是通过静态多态的方法调用了Derived这个子类定义的New函数来构造对象,同时把Derived这个子类中定义的Deleter这个删除器赋值给Object中定义的deleter_变量中,用于析构对象的时候用,下面是deleter_在Object中的定义:
class TVM_DLL Object {
public:
typedef void (*FDeleter)(Object* self);
protected:
FDeleter deleter_ = nullptr;
};
三、SimpleObjAllocator
SimpleObjAllocator中主要实现了基类中调用的New和Deleter函数,先看使用make_object构造单个Object时调用的New的实现:
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
template <typename T> class Handler {
using StorageType = typename std::aligned_storage<sizeof(T), alignof(T)>::type;
template <typename... Args>
static T* New(SimpleObjAllocator*, Args&&... args) {
StorageType* data = new StorageType();
new (data) T(std::forward<Args>(args)...);
return reinterpret_cast<T*>(data);
}
};
};
看到这里,所有的一切就真相大白了,make_object调用的New通过标准库的new来分配空间,然后再通过placement new来在分配的空间上构造对象,这里可能大家会有个疑问,就是为什么要用placement new来单独构造对象,TVM code base里有下面这段解释,我没做实验验证这个解释,大家有做过实验的可以留言告知一下:
// NOTE2: Use placement new to allocate
// This is used to get rid of warning when deleting a virtual
// class with non-virtual destructor.
再看make_object中给Object的deleter_变量赋值时调用的子类的Deleter函数的定义:
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
template <typename T> class Handler {
static Object::FDeleter Deleter() { return Deleter_; }
static void Deleter_(Object* objptr) {
T* tptr = static_cast<T*>(objptr);
tptr->T::~T();
delete reinterpret_cast<StorageType*>(tptr);
}
};
};
上面代码可以看到,删除器中先调用了析构函数,然后使用标准库的delete来释放内存。关于new/delete和inplacement new/delete,前文《内存管理:new and delete》有详细介绍,大家感兴趣可以点进去看看。
上面是make_object使用的New和Deleter,make_inplace_array所使用的New和Deleter的实现原理相同,只不过在分配和释放内存的时候,调用的是标准库的operator new/delete的数组版本,下面是make_inplace_array所使用的New和Deleter的关键代码:
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
template <typename ArrayType, typename ElemType> class ArrayHandler {
using StorageType = typename std::aligned_storage<sizeof(ArrayType), alignof(ArrayType)>::type;
template <typename... Args>
static ArrayType* New(SimpleObjAllocator*, size_t num_elems, Args&&... args) {
size_t unit = sizeof(StorageType);
size_t requested_size = num_elems * sizeof(ElemType) + sizeof(ArrayType);
size_t num_storage_slots = (requested_size + unit - 1) / unit;
StorageType* data = new StorageType[num_storage_slots];
new (data) ArrayType(std::forward<Args>(args)...);
return reinterpret_cast<ArrayType*>(data);
}
static Object::FDeleter Deleter() { return Deleter_; }
static void Deleter_(Object* objptr) {
ArrayType* tptr = static_cast<ArrayType*>(objptr);
tptr->ArrayType::~ArrayType();
StorageType* p = reinterpret_cast<StorageType*>(tptr);
delete[] p;
}
};
};
make_inplace_array所使用的New中有一点需要注意,它所申请的空间构成是下面这样子的:
图1
从图1可以看出,make_inplace_array申请的内存开头是Array作为header,后面紧跟着n个elements的空间,申请的时候把sizeof(array)+n*sizeof(element)+sizeof(padding)的空间作为num_storage_slots*sizeof(array)来申请,然后使用placement new来构造array header。
对于Deleter_函数,也有一点尤其需要注意,里面只调用了数组类型对象的析构函数,并没有调用里面所放元素的析构函数,需要user自己来手动在自己实现的Array中来调用存放元素的析构函数。Deleter_函数最后释放了数组对象和所有元素所占的内存空间。
四、summary and refenrence
TVM的内存分配器的实现很简单,但是它是整个TVM目前为Object分配对象的基石,非常重要值得一说,总结起来主要有下面几点:
使用基类+继承类,可以非常方便的扩展新的内存分配器,目前只实现了一个SimpleObjAllocator
使用CRTP实现静态多态,理论上来讲提高了内存分配的效率,因为如果使用动态多态的话,显然虚函数的调用开销要大于正常函数
既支持构造单个Object,也支持构造Object数组
本文参考的资料有下面这些:
TVM code base:https://github.com/apache/tvm/